import cv2  
import numpy as np  
import time  
  
class ImageProcessor:  
    def __init__(self, width=1920, height=1080):  
        self.width = width  
        self.height = height  
          
    def img_detect(self, BGR_image):  
        resized_image = self.resize(BGR_image, self.width, self.height)  
        return self.image_preprocess(resized_image)  
    
    def image_preprocess(self, BGR_image):  
        red_hsv_range = [(156, 43, 46), (180, 255, 255)]  
        mask = self.hsv_mask(BGR_image, red_hsv_range)  
        mask = self.close_operation(mask, 4)  
        contours = self.contour_process(BGR_image, mask)  
        # cv2.drawContours(BGR_image, contours, -1, (0, 0, 255), 3) # draw contours
        dst = cv2.bitwise_and(BGR_image, BGR_image, mask=mask)
        split_img = self.sort_and_display_contours(dst, contours) 
        return split_img  
      
    def hsv_mask(self, BGR_image, range):  
        HSV_image = cv2.cvtColor(BGR_image, cv2.COLOR_BGR2HSV)  
        return cv2.inRange(HSV_image, range[0], range[1])  
      
    def close_operation(self, image, iter):  
        kernel = np.ones((5,5),np.uint8) 
        return cv2.morphologyEx(image, cv2.MORPH_CLOSE, kernel, iterations=iter)  
      
    def contour_process(self, BGR_image, Bin_image):  
        contours, hierarchy = cv2.findContours(Bin_image, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)  
        return contours 
      
    def get_channel(self, image):  
        if image.ndim == 2:  
            channel = 1  
        else:  
            channel = image.shape[2]  
        print("Image channel = {}".format(channel))  
          
    def resize(self, image, width=None, height=None, inter=cv2.INTER_AREA):  
        dim = None  
        h, w, _ = image.shape  
        if width is None and height is None:  
            return image  
        if width is None:  
            r = height / float(h)  
            dim = (int(w * r), height)  
        else:  
            r = width / float(w)  
            dim = (width, int(h * r))  
        resized = cv2.resize(image, dim, interpolation=inter)  
        return resized  
      
    def sort_and_display_contours(self, image, contours):  
        # 计算每个轮廓的中心点  
        centers = [None] * len(contours)  
        for i, contour in enumerate(contours):  
            M = cv2.moments(contour)  
            if M["m00"] != 0:  
                cX = int(M["m10"] / M["m00"])  
                cY = int(M["m01"] / M["m00"])  
            else:  
                cX, cY = 0, 0  
            centers[i] = (cX, cY)  
    
        # 根据中心点的y坐标进行排序  
        contours = [x for _, x in sorted(zip(centers, contours), key=lambda pair: pair[0][1])]  
    
        # 计算输出图像的大小  
        # max_width = max(cv2.boundingRect(contour)[2] for contour in contours)
        # max_height = sum(cv2.boundingRect(contour)[3] for contour in contours)  
        # output_image = np.zeros((max_height, max_width, 3), dtype=np.uint8)  

        output_image = []
    
        # 绘制轮廓内的图像到输出图像  
        current_y = 0  
        for contour in contours:  
            x, y, w, h = cv2.boundingRect(contour)  
            if h < 10 or w <10:
                continue
            roi = image[y:y+h, x:x+w]  
            # output_image[current_y:current_y+h, :w] = roi 
            # current_y += h  
            output_image.append(roi) 
    
        # 返回排序后的轮廓图像  
        return output_image  
  
# 使用示例  
# processor = ImageProcessor(width=1920, height=1080)  
# processed_image = processor.img_detect(your_image)  # 替换your_image为你的图像数据


def test():
    img_path = r'image/raw_image/train_data.png'

    BGR_image = cv2.imread(img_path)

    processor = ImageProcessor(width=1920, height=1080)  
    processed_image = processor.img_detect(BGR_image)

    # # get_channel(image)

    print(len(processed_image))

    while True:
        cv2.imshow('image', processed_image[0])  # display image
        key = cv2.waitKey(24)
        if key == ord('q'):
            break
    


if __name__ == '__main__':
    test()
